import os
import random

trainval_percent = 0.2
train_percent = 0.8
# 读取xml路径
xmlfilepath = '../images/'
OS_LISTDIR = os.listdir(xmlfilepath)
# for i in range(len(OS_LISTDIR)):
#     OS_LISTDIR[i] = 'image/' + OS_LISTDIR[i]
total_xml = OS_LISTDIR

num = len(total_xml)
list = range(num)
# 分配各个数据的数量
tv = int(num * trainval_percent)  # 总数的20%
tr = int(tv * train_percent)  # 总数的20%的80%
trainval = random.sample(list, tv)  # 20%
train = random.sample(trainval, tr)  # 20%的80%

txtsavepath = '../txt/'
if not os.path.exists(txtsavepath):
    os.mkdir(txtsavepath)

ftrainval = open(txtsavepath + 'trainval.txt', 'w')
ftest = open(txtsavepath + 'test.txt', 'w')
ftrain = open(txtsavepath + 'train.txt', 'w')
fval = open(txtsavepath + 'val.txt', 'w')

for i in list:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:  # 20%
        ftrainval.write(name)  # 20%
        if i in train:
            ftest.write(name)  # 16%
        else:
            fval.write(name)  # 4%
    else:
        ftrain.write(name)  # 80%
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()
